package com.ruoyi.laiqianba.manager;

import com.ruoyi.db.domain.SseEvent;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class SseConnectionManager {
    
    private static final Map<String, SseEmitter> emitters = new ConcurrentHashMap<>();
    
    /**
     * 添加SSE连接
     */
    public SseEmitter addConnection(String clientId) {
        // 设置超时时间30分钟
        SseEmitter emitter = new SseEmitter(30 * 60 * 1000L);
        emitters.put(clientId, emitter);
        
        // 设置超时和错误处理
        emitter.onTimeout(() -> {
            emitter.complete();
            emitters.remove(clientId);
        });
        
        emitter.onCompletion(() -> {
            emitters.remove(clientId);
        });
        
        emitter.onError((ex) -> {
            emitter.complete();
            emitters.remove(clientId);
        });
        
        return emitter;
    }
    
    /**
     * 发送消息给指定客户端
     */
    public void sendToClient(String clientId, SseEvent event) {
        SseEmitter emitter = emitters.get(clientId);
        if (emitter != null) {
            try {
                emitter.send(SseEmitter.event()
                        .id(event.getEventId())
                        .name(event.getEventType())
                        .data(event)
                        .reconnectTime(5000L));
            } catch (Exception e) {
                emitter.completeWithError(e);
                emitters.remove(clientId);
            }
        }
    }
    
    /**
     * 广播消息给所有客户端
     */
    public void broadcast(SseEvent event) {
        List<String> disconnectedClients = new ArrayList<>();
        
        emitters.forEach((clientId, emitter) -> {
            try {
                emitter.send(SseEmitter.event()
                        .id(event.getEventId())
                        .name(event.getEventType())
                        .data(event)
                        .reconnectTime(5000L));
            } catch (Exception e) {
                disconnectedClients.add(clientId);
            }
        });
        
        // 清理断开连接的客户端
        disconnectedClients.forEach(emitters::remove);
    }
    
    /**
     * 移除连接
     */
    public void removeConnection(String clientId) {
        SseEmitter emitter = emitters.get(clientId);
        if (emitter != null) {
            emitter.complete();
            emitters.remove(clientId);
        }
    }
    
    /**
     * 获取活跃连接数
     */
    public int getActiveConnectionCount() {
        return emitters.size();
    }
}